package simulation;

import java.util.ArrayList;
import java.util.List;

import javax.vecmath.Matrix3d;
import javax.vecmath.Vector3d;

import rigidbody.RigidBody;
import simulation.ProjectedGaussSeidel.PGSResult;
import cern.colt.matrix.DoubleMatrix1D;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.impl.DenseDoubleMatrix1D;
import cern.colt.matrix.impl.DenseDoubleMatrix2D;
import cern.colt.matrix.impl.SparseDoubleMatrix2D;
import cern.colt.matrix.linalg.EigenvalueDecomposition;
import cern.jet.math.Functions;
import collision.Collision;

/**
 * Build the constraints as velocity-based complementarity formulation. 
 * See "Velocity-Based Shock Propagation for Multibody Dynamics Animation", Erleben, K., section 2.
 */
public class ProjectedGaussSeidelConstraintSolver implements ConstraintSolver {
	private static int MAX_ITERATION = 100;
	private static double TORLERANCE = 1e-5;
	private DoubleMatrix2D J, Minv, Jt, A;
	private DoubleMatrix1D f, u, b, p;
	private List<RigidBody> bodies;
	private List<Collision> collisions;
	private int k; // collisions.size()
	private int n; // bodies.size()
	private ProjectedGaussSeidel pgs;
	private double timestep;
	
	public ProjectedGaussSeidelConstraintSolver()
	{
		pgs = new ProjectedGaussSeidel();
	}

	@Override
	public boolean solveVelocities(Simulator sim)
	{
		timestep = sim.getTimestep();
		bodies = sim.getBodies();
		n = bodies.size();
		
		solveDynamics();
		
		// Integrate positions.
		for (int i = 0; i < n; i++)
			bodies.get(i).integratePosition(timestep);
		
//		solveProjection();

		// Clear all the effects from forces.
		for (int i = 0; i < n; i++)
			bodies.get(i).clearAccumulatedVars();
		
		return true;
	}
	
	/**
	 * Solve velocities to project bodies so that they do not penetrate.
	 * This method does not change positions.
	 */
	public void solveProjection()
	{
		// Get all pairs of collisions, despite of separating or not.
		getAllCollisions(false);
		if (collisions.size() == 0)
		{
			// No collisions, thus we don't need to do any correction.
			return;
		}
		
		solveProjectionWithCollisions();
	}

	public void solveProjectionWithCollisions() {
		// We've got some collisions, solve the velocities to eliminate them.
		buildJ();
		buildMInv();
		buildPenetration();
		// impulses = delta_t^2 * lambda
		DoubleMatrix1D impulses = solveImpulsesForConstraints(J, Minv, p);
		impulses.assign(Functions.div(timestep));
		
		// Calculate velocity delta from impulse.
		DoubleMatrix2D MinvJt = new SparseDoubleMatrix2D(6 * n, k);
		MinvJt.ensureCapacity(12 * k);
		Minv.zMult(Jt, MinvJt);
		DoubleMatrix1D du = new DenseDoubleMatrix1D(6 * n);
		MinvJt.zMult(impulses, du);

		// Apply only the calculated velocity to push objects apart.
		
		// Backup and reset velocities.
		for (int i = 0; i < n; i++)
		{
			bodies.get(i).backupVelocity();
			bodies.get(i).setAngularMomentum(new Vector3d());
			bodies.get(i).setLinearMomentum(new Vector3d());
		}
		
		// Calculate linear and angular momentum from velocity.
		for (int i = 0; i < n; i++)
		{
			int s = i * 6;
			Vector3d lin = new Vector3d(du.get(s), du.get(s + 1), du.get(s + 2));
			s += 3;
			Vector3d ang = new Vector3d(du.get(s), du.get(s + 1), du.get(s + 2));
			bodies.get(i).addMomentumFromVelocityDelta(lin, ang);
		}
		
		// Correct positions.
		for (int i = 0; i < n; i++)
		{
			bodies.get(i).integratePosition(timestep);
		}
		
		// Restore velocities.
		for (int i = 0; i < n; i++)
			bodies.get(i).restoreVelocity();
	}
	
	private void buildPenetration()
	{
		p = new DenseDoubleMatrix1D(k);
		for (int i = 0; i < k; i++)
			p.set(i, -collisions.get(i).penetrationDepth);
	}
	
	/**
	 * Solve velocities by incorporating forces.
	 * This method does not change positions.
	 */
	public void solveDynamics()
	{
		getAllCollisions(false);
		
		if (collisions.size() == 0)
		{
			// No collisions, let's integrate the velocities.
			for (int i = 0; i < n; i++)
				bodies.get(i).integrateVelocity(timestep);
			return;
		}
		
//		getAllCollisions(false);
		
		solveDynamicsWithCollisions();
	}

	public void solveDynamicsWithCollisions() {
		// We've got some collisions, solve the velocities to eliminate them.
		buildJ();
		buildMInv();
		buildF();
		buildB();
		DoubleMatrix1D impulses = solveImpulsesForConstraints(J, Minv, b);
		
		// Integrate external forces.
		for (int i = 0; i < n; i++)
			bodies.get(i).integrateVelocity(timestep);
		
		// Calculate velocity delta from impulse.
		DoubleMatrix2D MinvJt = new SparseDoubleMatrix2D(6 * n, k);
		MinvJt.ensureCapacity(12 * k);
		Minv.zMult(Jt, MinvJt);
		DoubleMatrix1D du = new DenseDoubleMatrix1D(6 * n);
		MinvJt.zMult(impulses, du);
		
		// Calculate linear and angular momentum from velocity.
		for (int i = 0; i < n; i++)
		{
			int s = i * 6;
			Vector3d lin = new Vector3d(du.get(s), du.get(s + 1), du.get(s + 2));
			s += 3;
			Vector3d ang = new Vector3d(du.get(s), du.get(s + 1), du.get(s + 2));
			bodies.get(i).addMomentumFromVelocityDelta(lin, ang);
		}
	}
	
	public DoubleMatrix1D solveImpulsesForConstraints(DoubleMatrix2D J,
			DoubleMatrix2D Minv, DoubleMatrix1D b)
	{
		buildA();
		
		DoubleMatrix1D lo = new DenseDoubleMatrix1D(k);
		DoubleMatrix1D hi = new DenseDoubleMatrix1D(k);
		lo.assign(0);
		hi.assign(Double.MAX_VALUE);
		
		PGSResult result = pgs.solve(A, b, lo, hi,
				MAX_ITERATION, TORLERANCE);
		if (result.iterationProcessed >= MAX_ITERATION)
		{
			System.out.println("Solver reached max iteration.");
		}
		return result.result;
	}
	
	private void buildA()
	{
		Jt = J.viewDice();
		
		// A = J * M^-1 * J^T
		DoubleMatrix2D JMinv = new SparseDoubleMatrix2D(k, 6 * n);
		J.zMult(Minv, JMinv);
		A = new SparseDoubleMatrix2D(k, k);
		A.ensureCapacity(12 * k);
		JMinv.zMult(Jt, A);
		
		reviseA();
	}
	
	private void reviseA()
	{
		EigenvalueDecomposition ed = new EigenvalueDecomposition(A);
		DoubleMatrix2D d = ed.getD();
		DoubleMatrix2D v = ed.getV();
		DoubleMatrix2D vt = v.viewDice();
		for (int i = 0; i < d.rows(); i++)
		{
			if (d.get(i, i) < 0.1)
				d.set(i, i, 0.1);
		}
		
		DoubleMatrix2D vd = new SparseDoubleMatrix2D(k, k);
		v.zMult(d, vd);
		vd.zMult(vt, A);
	}
	
	private void buildB()
	{
		u = new DenseDoubleMatrix1D(6 * n);
		for (int i = 0; i < n; i++)
		{
			RigidBody b = bodies.get(i);
			assignVectorToVector(u, b.getVelocity(), i * 6);
			// Angular velocity should have been calculated in buildF().
			assignVectorToVector(u, b.getAngularVelocity(), i * 6 + 3);
		}
		// uext = u + timestep * M^-1 * f.
		DoubleMatrix1D uext = new DenseDoubleMatrix1D(6 * n);
		Minv.zMult(f, uext);
		uext.assign(Functions.mult(timestep));
		uext.assign(u, Functions.plus);
		
		// b = J * (u + timestep * M^-1 * f)
		b = new DenseDoubleMatrix1D(k);
		J.zMult(uext, b);
	}
	
	private void buildJ()
	{
		J = new SparseDoubleMatrix2D(k, n * 6);
		J.ensureCapacity(k * 12);
		for (int i = 0; i < collisions.size(); i++)
		{
			Collision c = collisions.get(i);
			RigidBody b1 = c.object1;
			RigidBody b2 = c.object2;
			Vector3d normal = new Vector3d(c.normal);
			normal.negate(); // Normal definition in Collision is the negative direction defined here.
			
			Vector3d r1 = new Vector3d(c.positionInWorld);
			r1.sub(b1.getPosition());
			Vector3d r2 = new Vector3d(c.positionInWorld);
			r2.sub(b2.getPosition());
			
			Vector3d rxn1 = new Vector3d();
			rxn1.cross(r1, normal);
			Vector3d rxn2 = new Vector3d();
			rxn2.cross(r2, normal);
			
			// J_{i, b1*6-b1*6+5} = [n^T, (r1 x n)^T]
			assignVectorToMatrix(J, normal, i, 6 * b1.getId(), true, false);
			assignVectorToMatrix(J, rxn1, i, 6 * b1.getId() + 3, true, false);

			// J_{i, b2*6-b2*6+5} = [-n^T, -(r2 x n)^T]			
			assignVectorToMatrix(J, normal, i, 6 * b2.getId(), true, true);
			assignVectorToMatrix(J, rxn2, i, 6 * b2.getId() + 3, true, true);
		}
	}
	
	private void buildMInv()
	{
		Minv = new SparseDoubleMatrix2D(n * 6, n * 6);
		Minv.ensureCapacity(12 * n);
		for (int i = 0; i < n; i++)
		{
			int start = i * 6;
			double massinv = bodies.get(i).getMassInv();
			for (int j = 0; j < 3; j++)
			{
				Minv.setQuick(start + j, start + j, massinv);
			}
			RigidBody b = bodies.get(i);
			
			Matrix3d inertiaInv = b.calculateInertiaInverse();
			start += 3;
			for (int x = 0; x < 3; x++)
				for (int y = 0; y < 3; y++)
				{
					Minv.setQuick(start + x, start + y, inertiaInv.getElement(x, y));
				}
		}
	}
	
	private void buildF()
	{
		f = new DenseDoubleMatrix1D(n * 6);
		f.ensureCapacity(6 * n);
		for (int i = 0; i < n; i++)
		{
			RigidBody b = bodies.get(i);
			Vector3d flin = b.getAccumulatedForce();
			// f_ang = torque - ang_velocity x ang_momentum
			// See Siggraph 2001 rigidbody course, notes appendix C.2
			Vector3d fang = new Vector3d(b.getAccumulatedTorque());
			Vector3d l = b.getAngularMomentum();
			Vector3d w = b.calcAngularVelocity();
			Vector3d wxl = new Vector3d();
			wxl.cross(w, l);
			fang.sub(wxl);
			assignVectorToVector(f, flin, i * 6);
			assignVectorToVector(f, fang, i * 6 + 3);
		}
	}
	
	private void getAllCollisions(boolean nonSeparating)
	{
		collisions = new ArrayList<Collision>();
		for (int i = 0; i < bodies.size(); i++)
		{
			RigidBody b1 = bodies.get(i);
			for (int j = 0; j < i; j++)
			{
				RigidBody b2 = bodies.get(j);
				collisions.addAll(RigidBody.getCollisionsBetween(b1, b2, nonSeparating));
			}
		}
		
		k = collisions.size();
	}
	
	private void assignVectorToMatrix(DoubleMatrix2D m, Vector3d v, int row, int col, boolean transpose, boolean negate)
	{
		double x = v.x, y = v.y, z = v.z;
		if (negate)
		{
			x = -x;
			y = -y;
			z = -z;
		}
		
		if (!transpose)
		{
			m.setQuick(row, col, x);
			m.setQuick(row + 1, col, y);
			m.setQuick(row + 2, col, z);
		}
		else
		{
			m.setQuick(row, col, x);
			m.setQuick(row, col + 1, y);
			m.setQuick(row, col + 2, z);
		}
	}
	
	private void assignVectorToVector(DoubleMatrix1D m, Vector3d v, int pos)
	{
		m.setQuick(pos, v.x);
		m.setQuick(pos + 1, v.y);
		m.setQuick(pos + 2, v.z);
	}
	
	@Override
	public void rigidBodyAdded(Simulator sim)
	{
	}

	/**
	 * Method for testing.
	 */
	public void setTimestep(double timestep)
	{
		this.timestep = timestep;
	}
	
	/**
	 * Method for testing.
	 */
	public void setBodies(List<RigidBody> bodies)
	{
		this.bodies = bodies;
		n = bodies.size();
	}
	
	/**
	 * Method for testing.
	 */
	public void setCollisions(List<Collision> collisions)
	{
		this.collisions = collisions;
		k = collisions.size();
	}
}
